# -*- coding: utf-8 -*-

import logger.logger as logger
from sklearn.externals import joblib
from sklearn.naive_bayes import GaussianNB


class NaiveBayes():
    def __init__(self):
        self.model = None
        self.parameter = {}

    def train(self, x_train, y_train, x_test, y_test, model_path):
        logger.info("training...")
        self.model = GaussianNB(**self.parameter)
        self.model.fit(x_train, y_train)
        joblib.dump(self.model, model_path)
        acc = self.model.score(x_test, y_test)
        logger.info("acc: {}".format(str(acc)))

    def cross_validation(self, train_x, train_y, test_vecs, y_test, model_path, param_grid=None):
        from sklearn.model_selection import GridSearchCV
        model = GaussianNB()
        param_grid = [
            {'kernel': ['rbf'],
             'C': [0.001, 0.01, 0.1, 1, 10, 100],
             'gamma': [0.001, 0.01, 0.1, 0.2, 0.4, 0.6, 0.8, 1, 1.2, 1.4, 1.6]},

            {'kernel': ['linear'],
             'C': [0.001, 0.01, 0.1, 1, 10, 10]}
        ]
        grid_search = GridSearchCV(model, param_grid, n_jobs=4, verbose=1, cv=5)
        grid_search.fit(train_x, train_y)
        best_parameters = grid_search.best_estimator_.get_params()
        for para, val in list(best_parameters.items()):
            print(para, val)

        self.model = GaussianNB(kernel='rbf', C=best_parameters['C'], gamma=best_parameters['gamma'], probability=True)
        self.model.fit(train_x, train_y)
        joblib.dump(self.model, model_path)
        acc = self.model.score(test_vecs, y_test)
        logger.info("acc: {}".format(str(acc)))

    def load_model(self, model_path):
        self.model = joblib.load(model_path)

    def predict(self, x_vecs):
        if not self.model:
            logger.info("please load naive bayes model first!!!")
            return
        y_pred = self.model.predict(x_vecs)
        return y_pred


if __name__ == '__main__':
    pass
